/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file edge_sub.cpp
 */
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelEdgeSub {
public:
    __aicore__ inline KernelEdgeSub() {}
    __aicore__ inline void Init(GM_ADDR img1, GM_ADDR img2, GM_ADDR out,
                                uint32_t formerNum, uint32_t formerLength, uint32_t tailLength, uint32_t width, uint32_t heigh,
                                TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        
        outGm_original.SetGlobalBuffer((__gm__ uint8_t*)out);
        this->coreDataNum = 2101408;
        if(GetBlockIdx() == 7) 
        {
            this->coreDataNum = 2101440;
        }
        img1Gm.SetGlobalBuffer((__gm__ uint8_t*)img1 + 2101408 * GetBlockIdx(), this->coreDataNum);
        img2Gm.SetGlobalBuffer((__gm__ uint8_t*)img2 + 2101408 * GetBlockIdx(), this->coreDataNum);
        outGm.SetGlobalBuffer((__gm__ uint8_t*)out + 2101408 * GetBlockIdx(), this->coreDataNum);
        this->width = width;
        this->heigh = heigh;
        this->tileDataNum = 32256;
        this->tailDataNum = this->coreDataNum % this->tileDataNum;
        this->tileNum = this->coreDataNum / this->tileDataNum;
        if(this->tailDataNum == 0)
        {
            this->tileNum -= 1; 
            this->tailDataNum = this->tileDataNum;
        }
        pipeIn->InitBuffer(inQueueImg1, BUFFER_NUM, this->tileDataNum * 2 * sizeof(uint8_t));
        pipeIn->InitBuffer(outQueueOut, BUFFER_NUM, this->tileDataNum * sizeof(uint8_t));
        pipeIn->InitBuffer(QueueTmp1, this->tileDataNum * sizeof(half));
    }
    __aicore__ inline void Process() {
        int32_t loopCount = this->tileNum;
        this->processDataNum = this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
        this->processDataNum = this->tailDataNum;
        CopyIn(loopCount);
        Compute(loopCount);
        CopyOut(loopCount);
    }
private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<uint8_t> img1Local = inQueueImg1.AllocTensor<uint8_t>();
        DataCopy(img1Local, img1Gm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(img1Local[this->tileDataNum], img2Gm[progress * this->tileDataNum], this->processDataNum);
        inQueueImg1.EnQue(img1Local);
    }
    __aicore__ inline void Compute(int32_t progress) {
        LocalTensor<uint8_t> img1Local = inQueueImg1.DeQue<uint8_t>();
        LocalTensor<uint8_t> outLocal = outQueueOut.AllocTensor<uint8_t>();
        auto tmp1 = QueueTmp1.Get<half>();
        Cast(tmp1, img1Local, RoundMode::CAST_NONE, this->processDataNum);
        Cast(img1Local.ReinterpretCast<half>(), img1Local[this->tileDataNum], RoundMode::CAST_NONE, this->processDataNum);
        Adds(tmp1, tmp1, (half)(10), this->processDataNum);
        SubRelu(tmp1, img1Local.ReinterpretCast<half>(), tmp1, this->processDataNum);
        Cast(outLocal, tmp1, RoundMode::CAST_NONE, this->processDataNum);
        inQueueImg1.FreeTensor(img1Local);
        outQueueOut.EnQue<uint8_t>(outLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<uint8_t> outLocal = outQueueOut.DeQue<uint8_t>();
        DataCopy(outGm[progress * this->tileDataNum], outLocal, this->processDataNum);
        outQueueOut.FreeTensor(outLocal);
    }
private:
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueImg1,inQueueImg2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOut;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2;

    GlobalTensor<uint8_t> img1Gm, img2Gm;
    GlobalTensor<uint8_t> outGm, outGm_original;

    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;

    uint32_t width;
    uint32_t heigh;
};
extern "C" __global__ __aicore__ void edge_sub(GM_ADDR img1, GM_ADDR img2, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    KernelEdgeSub op;
    op.Init(img1, img2, out,
            tiling_data.formerNum, tiling_data.formerLength, tiling_data.tailLength, tiling_data.width, tiling_data.heigh, 
            &pipe);  
    op.Process();
}